# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import os

from coin.base.datetime_util import to_datetime, to_timestamp_int
from coin.base.param_util import to_list, to_tuple
from coin.strategy.mm.subscription import SubscriptionRequest, get_feed_sub_request
from coin.feed.fastfeed.util import split_time_range
from coin.feed.fastfeed.constants import DEFAULT_FEED_CACHE_DIR
from cc.coin1.fastfeed.pyfastfeed import _ProtoFeedBuilder, _FeedCacheRunner


def is_native_callback(callback):
  try:
    return callback.native
  except AttributeError:
    return False


class ProtoFeedBuilder:
  def __init__(self, *, sub_req=None, builder=None):
    if builder:
      self._builder = builder
    else:
      assert sub_req
      self._builder = _ProtoFeedBuilder(sub_req.to_str())

  @property
  def cc_builder(self):
    return self._builder.cc_builder

  def update(self, serialized_msg):
    self._builder.update(serialized_msg)

  def subscribe(self, symbols, callback):
    self.subscribe_book(symbols, callback)

  def subscribe_book(self, symbols, callback):
    for symbol in to_list(symbols):
      # FIX: Subscribe both temporarily.
      # coin1 use full symbol and coin2 use just symbol.
      # symbol: "BTC-USD.20200327.Huobi.BTC200327",
      # symbol: "BTC-USD.20200228"
      self._builder.subscribe_book(symbol.full_symbol, callback)
      self._builder.subscribe_book(symbol.symbol, callback)

  def subscribe_trade(self, symbols, callback):
    for symbol in to_list(symbols):
      self._builder.subscribe_trade(symbol.full_symbol, callback)
      self._builder.subscribe_trade(symbol.symbol, callback)

  def subscribe_instrument(self, symbols, callback):
    for symbol in to_list(symbols):
      self._builder.subscribe_instrument(symbol.full_symbol, callback)
      self._builder.subscribe_instrument(symbol.symbol, callback)


def hack_sub_req_name(machine_chosen, curr_sub_req, sub_req_hack_map):
  if sub_req_hack_map is not None:
    return sub_req_hack_map.get(curr_sub_req, curr_sub_req).to_str()
  return curr_sub_req.to_str().replace("v5_swap", "v5-swap.realtime_trade_move2bp")


def _run_daily_from_feed_cache(*,
                               feed_sub_request,
                               on_book_reset_callback,
                               date,
                               from_ts,
                               to_ts,
                               machine,
                               feed_cache_dir,
                               worker_id,
                               sub_req_to_machine_override_map,
                               use_simplified_filename,
                               sub_req_hack_map=None):
  assert date == to_datetime(from_ts).date()
  assert isinstance(machine, str)

  feed_cache_dir = [os.path.expanduser(dir) for dir in to_tuple(feed_cache_dir)]
  runner = _FeedCacheRunner(date.strftime('%Y%m%d'), feed_cache_dir)
  sub_req_to_machine_override_map = sub_req_to_machine_override_map or dict()

  sub_req_hack_map_backward = None
  if sub_req_hack_map is not None:
    sub_req_hack_map_backward = {value: key for key, value in sub_req_hack_map.items()}

  def _on_book_reset(feed_cache_key, feed_builder):
    if feed_cache_key.sub_req.endswith("v5-swap.realtime"):
      sub_req = SubscriptionRequest.from_str(
          feed_cache_key.sub_req.replace("v5-swap.realtime", "v5_swap"))
    else:
      sub_req = SubscriptionRequest.from_str(feed_cache_key.sub_req)

    if sub_req_hack_map_backward is not None:
      sub_req = sub_req_hack_map_backward.get(sub_req, sub_req)

    for callback in to_list(on_book_reset_callback):
      try:
        callback(sub_req, ProtoFeedBuilder(builder=feed_builder), feed_cache_key=feed_cache_key)
      except TypeError:
        callback(sub_req, ProtoFeedBuilder(builder=feed_builder))

  runner.set_book_reset_callback(_on_book_reset)

  for sub_req, products in feed_sub_request.gen_iter():
    machine_chosen = sub_req_to_machine_override_map.get(sub_req, machine)
    if worker_id is None:
      if use_simplified_filename:
        runner.add_cache_key_with_feed_arb(
            machine_chosen,
            hack_sub_req_name(machine_chosen, sub_req, sub_req_hack_map),
            [p.symbol for p in products])
      else:
        runner.add_cache_key_with_feed_arb(
            machine_chosen,
            hack_sub_req_name(machine_chosen, sub_req, sub_req_hack_map),
            [str(p) for p in products])
    else:
      if use_simplified_filename:
        runner.add_cache_key(machine_chosen,
                             worker_id,
                             hack_sub_req_name(machine_chosen, sub_req, sub_req_hack_map),
                             [p.symbol for p in products])
      else:
        runner.add_cache_key(machine_chosen,
                             worker_id,
                             hack_sub_req_name(machine_chosen, sub_req, sub_req_hack_map),
                             [str(p) for p in products])

  runner.run(from_ts, to_ts)
  return runner


def run_from_feed_cache(feed_sub_request_or_products,
                        on_book_reset_callback,
                        from_ts,
                        to_ts,
                        *,
                        machine,
                        sub_req_hack_map=None,
                        feed_cache_dir=None,
                        worker_id=None,
                        sub_req_to_machine_override_map=None,
                        use_simplified_filename=True):
  feed_sub_request = get_feed_sub_request(feed_sub_request_or_products)
  feed_cache_dir = feed_cache_dir or DEFAULT_FEED_CACHE_DIR

  if isinstance(machine, list):
    assert len(machine) == 1
    machine = machine[0]

  if isinstance(worker_id, list):
    assert set(worker_id) == {'1', '2'}
    worker_id = None

  ranges = split_time_range(from_ts, to_ts)
  prev_runner = None
  for date, from_ts, to_ts in ranges:
    runner = _run_daily_from_feed_cache(
        feed_sub_request=feed_sub_request,
        on_book_reset_callback=on_book_reset_callback,
        date=date,
        from_ts=from_ts,
        to_ts=to_ts,
        sub_req_hack_map=sub_req_hack_map,
        machine=machine,
        feed_cache_dir=feed_cache_dir,
        worker_id=worker_id,
        sub_req_to_machine_override_map=sub_req_to_machine_override_map,
        use_simplified_filename=use_simplified_filename)
    # Keep the created fastfeed objects alive.
    prev_runner, runner = runner, None
